
# 新增以下关键导入
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

# 原有其他导入
import os
import base64
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import padding as asym_padding

def print_header(title):
    """美观的标题打印"""
    print(f"\n{'='*40}")
    print(f"=== {title}")
    print(f"{'='*40}\n")

def print_hex(data, prefix="", length=16):
    """以十六进制格式打印二进制数据"""
    hex_str = data.hex()[:length*2]
    truncated = "(truncated)" if len(data) > length else ""
    print(f"{prefix}{hex_str} {truncated}[Total {len(data)} bytes]")

# ----------------- 私钥加载阶段 -----------------
print_header("1. 加载虚拟机私钥")

try:
    with open("vm.key", "rb") as f:
        vm_private_key = serialization.load_pem_private_key(
            f.read(),
            password=None
        )
    print("[私钥加载成功]")
    print(f"  密钥类型: {type(vm_private_key).__name__}")
    print(f"  密钥长度: {vm_private_key.key_size} bits\n")

except FileNotFoundError:
    print("[错误] 找不到 vm.key 文件！")
    exit(1)
except Exception as e:
    print(f"[错误] 私钥加载失败: {str(e)}")
    exit(1)

# ----------------- 会话密钥解密 -----------------
print_header("2. 解密会话密钥")

try:
    # 假设从文件读取加密的会话密钥
    with open("encrypted_session_key.bin", "rb") as f:
        encrypted_key = f.read()
    
    print_hex(encrypted_key, "接收的加密密钥: ", length=12)
    
    decrypted_key = vm_private_key.decrypt(
        encrypted_key,
        asym_padding.OAEP(
            mgf=asym_padding.MGF1(algorithm=hashes.SHA256()),
            algorithm=hashes.SHA256(),
            label=None
        )
    )
    
    print("\n[解密成功]")
    print(f"  解密后长度: {len(decrypted_key)} bytes")
    print_hex(decrypted_key, "原始会话密钥: ", length=8)
    print(f"  Base64预览: {base64.b64encode(decrypted_key[:8]).decode()}...\n")

except Exception as e:
    print(f"[解密失败] 错误类型: {type(e).__name__}, 详情: {str(e)}")
    exit(1)

# ----------------- 数据加密阶段 -----------------
print_header("3. 加密敏感数据")

try:
    # 生成初始化向量
    iv = os.urandom(16)
    print_hex(iv, "生成的IV: ", length=8)
    
    # 加密示例数据
    plaintext = b"Secret Data: VM->Host Communication Test"
    cipher = Cipher(algorithms.AES(decrypted_key), modes.GCM(iv))
    encryptor = cipher.encryptor()
    encrypted_data = encryptor.update(plaintext) + encryptor.finalize()
    tag = encryptor.tag
    
    print("\n[加密结果]")
    print(f"  原始数据长度: {len(plaintext)} bytes")
    print(f"  加密后长度: {len(encrypted_data)} bytes")
    print_hex(encrypted_data, "加密数据: ", length=12)
    print_hex(tag, "认证标签: ", length=6)

except Exception as e:
    print(f"[加密失败] 错误类型: {type(e).__name__}, 详情: {str(e)}")
    exit(1)

# ----------------- 数据签名阶段 -----------------
print_header("4. 生成数字签名")

try:
    data_to_sign = iv + tag + encrypted_data
    print_hex(data_to_sign, "待签名数据组合: ", length=20)
    
    signature = vm_private_key.sign(
        data_to_sign,
        asym_padding.PSS(
            mgf=asym_padding.MGF1(hashes.SHA256()),
            salt_length=asym_padding.PSS.MAX_LENGTH
        ),
        hashes.SHA256()
    )
    
    print("\n[签名生成成功]")
    print(f"  签名长度: {len(signature)} bytes")
    print_hex(signature, "签名内容: ", length=12)
    print(f"  Base64预览: {base64.b64encode(signature[:16]).decode()}...")

except Exception as e:
    print(f"[签名失败] 错误类型: {type(e).__name__}, 详情: {str(e)}")
    exit(1)

# ----------------- 发送准备阶段 -----------------
print_header("5. 发送准备")

print("[即将发送的数据包]")
print(f"  IV长度: {len(iv)} bytes")
print(f"  Tag长度: {len(tag)} bytes")
print(f"  加密数据长度: {len(encrypted_data)} bytes")
print(f"  签名长度: {len(signature)} bytes")
print("\n[下一步操作] 请将以下文件发送到物理机:")
print("  - iv.bin (16 bytes)")
print("  - tag.bin (16 bytes)")
print("  - encrypted_data.bin")
print("  - signature.bin")
print("  - vm.crt (证书文件)")

# 保存发送数据到文件（示例）
with open("iv.bin", "wb") as f: f.write(iv)
with open("tag.bin", "wb") as f: f.write(tag)
with open("encrypted_data.bin", "wb") as f: f.write(encrypted_data)
with open("signature.bin", "wb") as f: f.write(signature)

print("\n" + "="*40)
print("=== 虚拟机端操作完成 ===")
print("="*40)